package ml.features;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.TreeMap;
import java.util.TreeSet;

import corpus.SentPolarityCorpus;

import ml.Data;

import parser.ParserAnnotation;
import parser.ParserPipeline;

public class POSExtractor extends FeatureExtractor {
	private String[] poss;

	public POSExtractor(int offset, SentPolarityCorpus corpus, ParserPipeline p) {
		super(offset);
		this.poss = getBagOfPOSArray(corpus, p);
	}

	/**
	 * Cached version.
	 */
	public POSExtractor(int offset) throws IOException {
		super(offset);
		poss = loadList(new File("features/pos"));
	}

	private Set<String> getBagOfPOS(SentPolarityCorpus corpus, ParserPipeline p) {
		TreeSet<String> set = new TreeSet<String>();
		for (String sent : corpus.getCorpus().keySet()) {
			ParserAnnotation pa = corpus.getParsedPA(sent, p);
			for (String pos : pa.getPOS(0))
				set.add(pos);
		}
		return set;
	}

	private String[] getBagOfPOSArray(SentPolarityCorpus corpus,
			ParserPipeline p) {
		Set<String> set = getBagOfPOS(corpus, p);
		String[] pos = new String[set.size()];
		return set.toArray(pos);
	}

	public List<Data> extract(ParserAnnotation pa, int sentence) {
		List<Data> datas = new ArrayList<Data>(pa.getNbrSentences());
		List<String> sentPOS = pa.getPOS(sentence);
		Map<String, Integer> posMap = new TreeMap<String, Integer>();
		int nbrTokens = pa.getTokens(sentence).size();
		for (String s : sentPOS) {
			if (posMap.containsKey(s))
				posMap.put(s, posMap.get(s) + 1);
			else
				posMap.put(s, 1);
		}
		for (Entry<String, Integer> e : posMap.entrySet()) {
			int index = Arrays.binarySearch(poss, e.getKey());
			datas.add(new Data((double) e.getValue() / nbrTokens, index + offset));
		}
		return datas;
	}

	public List<Data>[] extractSequence(ParserAnnotation pa, int sentence) {
		List<Data>[] datas = new List[nbrColumns()];
		for (int i = 0; i < datas.length; i++)
			datas[i] = new ArrayList();
		List<String> sentPos = pa.getPOS(sentence);

		for (String tokenPos : sentPos) {
			for (int i = 0; i < poss.length; i++) {
				String pos = poss[i];
				datas[i].add(new Data(pos.equals(tokenPos), i));
			}
		}
		return datas;
	}

	public int nbrColumns() {
		return poss.length;
	}

	@Override
	public String[] wekaDescription() {
		String[] description = new String[poss.length];
		for (int i = 0; i < description.length; i++) {
			description[i] = "\"POS " + poss[i].replace("\"", "\\\\")
					+ "\" numeric";
		}
		return description;
	}

	@Override
	public void store() throws IOException {
		File out = new File("features/pos");
		storeSet(Arrays.asList(poss), out);
	}
}
